#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
优化版 Mermaid to Draw.io 转换器
修复：XML 压缩、语法错误、布局问题
"""

import os
import re
import xml.etree.ElementTree as ET
from xml.dom import minidom
import zlib
import base64


class SimpleMermaidConverter:
    def __init__(self):
        self.node_id_counter = 0
        self.edge_id_counter = 0

    def generate_node_id(self):
        self.node_id_counter += 1
        return f"node_{self.node_id_counter}"

    def generate_edge_id(self):
        self.edge_id_counter += 1
        return f"edge_{self.edge_id_counter}"

    def create_drawio_skeleton(self, title="Diagram"):
        """创建 draw.io 的骨架结构（不包含 diagram 内容）"""
        mxfile = ET.Element("mxfile")
        mxfile.set("host", "app.diagrams.net")
        mxfile.set("modified", "2024-01-01T00:00:00.000Z")
        mxfile.set("agent", "5.0")
        mxfile.set("version", "22.1.16")
        mxfile.set("type", "device")

        # diagram 元素（内容稍后填充）
        diagram = ET.SubElement(mxfile, "diagram")
        diagram.set("id", "diagram_1")
        diagram.set("name", title)

        return mxfile, diagram

    def create_mxgraphmodel_xml(self, nodes, edges, subgraphs, positions):
        """构建 mxGraphModel XML（明文）"""
        mxGraphModel = ET.Element("mxGraphModel")
        mxGraphModel.set("dx", "1422")
        mxGraphModel.set("dy", "794")
        mxGraphModel.set("grid", "1")
        mxGraphModel.set("gridSize", "10")
        mxGraphModel.set("guides", "1")
        mxGraphModel.set("tooltips", "1")
        mxGraphModel.set("connect", "1")
        mxGraphModel.set("arrows", "1")
        mxGraphModel.set("fold", "1")
        mxGraphModel.set("page", "1")
        mxGraphModel.set("pageScale", "1")
        mxGraphModel.set("pageWidth", "1169")
        mxGraphModel.set("pageHeight", "827")
        mxGraphModel.set("math", "0")
        mxGraphModel.set("shadow", "0")

        root = ET.SubElement(mxGraphModel, "root")

        # 添加根 cell
        cell0 = ET.SubElement(root, "mxCell")
        cell0.set("id", "0")

        cell1 = ET.SubElement(root, "mxCell")
        cell1.set("id", "1")
        cell1.set("parent", "0")

        # 创建子图
        for subgraph in subgraphs:
            pos = positions.get(subgraph['id'], {'x': 50, 'y': 50, 'width': 300, 'height': 250})
            self.create_subgraph_element(root, subgraph, pos)

        # 创建节点
        for node in nodes:
            pos = positions.get(node['id'], {'x': 50, 'y': 50, 'width': 120, 'height': 60})
            self.create_node_element(root, node, pos)

        # 创建边
        for edge in edges:
            if edge['source'] in [n['id'] for n in nodes] and edge['target'] in [n['id'] for n in nodes]:
                self.create_edge_element(root, edge)

        return mxGraphModel

    def create_node_element(self, root, node, pos):
        x, y, width, height = pos['x'], pos['y'], pos['width'], pos['height']
        cell = ET.SubElement(root, "mxCell")
        cell.set("id", node['id'])
        cell.set("value", node['label'])
        cell.set("style", node['style'])
        cell.set("vertex", "1")
        cell.set("parent", "1")

        geometry = ET.SubElement(cell, "mxGeometry")
        geometry.set("x", str(x))
        geometry.set("y", str(y))
        geometry.set("width", str(width))
        geometry.set("height", str(height))
        geometry.set("as", "geometry")

    def create_edge_element(self, root, edge):
        cell = ET.SubElement(root, "mxCell")
        cell.set("id", self.generate_edge_id())
        cell.set("value", edge['label'])
        cell.set("style", edge['style'])
        cell.set("edge", "1")
        cell.set("parent", "1")
        cell.set("source", edge['source'])
        cell.set("target", edge['target'])

        geometry = ET.SubElement(cell, "mxGeometry")
        geometry.set("relative", "1")
        geometry.set("as", "geometry")

    def create_subgraph_element(self, root, subgraph, pos):
        x, y, width, height = pos['x'], pos['y'], pos['width'], pos['height']
        cell = ET.SubElement(root, "mxCell")
        cell.set("id", subgraph['id'])
        cell.set("value", subgraph['name'])
        cell.set("style", "swimlane;fontStyle=0;childLayout=stackLayout;horizontal=1;startSize=30;"
                             "horizontalStack=0;resizeParent=1;resizeParentMax=0;resizeLast=0;"
                             "collapsible=1;marginBottom=0;whiteSpace=wrap;html=1;fillColor=#f5f5f5;"
                             "strokeColor=#666666;fontColor=#333333;")
        cell.set("vertex", "1")
        cell.set("parent", "1")

        geometry = ET.SubElement(cell, "mxGeometry")
        geometry.set("x", str(x))
        geometry.set("y", str(y))
        geometry.set("width", str(width))
        geometry.set("height", str(height))
        geometry.set("as", "geometry")

    def compress_xml(self, element):
        """将 XML 元素压缩为 base64(deflate(bytes))"""
        raw_xml = ET.tostring(element, encoding='utf-8')
        compressed = zlib.compress(raw_xml, 9)
        return base64.b64encode(compressed).decode('utf-8')

    def parse_mermaid_code(self, mermaid_code):
        lines = mermaid_code.strip().split('\n')
        nodes = []
        edges = []
        subgraphs = []
        current_subgraph = None

        for line in lines:
            line = line.strip()
            if not line or line.startswith('```'):
                continue
            if line.startswith(('graph', 'flowchart', 'sequenceDiagram', 'stateDiagram')):
                continue

            # 子图
            if line.startswith('subgraph'):
                match = re.match(r'subgraph\s+"([^"]+)"', line)
                if match:
                    current_subgraph = {
                        'id': f"subgraph_{len(subgraphs)}",
                        'name': match.group(1),
                        'nodes': []
                    }
                    subgraphs.append(current_subgraph)
                continue

            if line == 'end' and current_subgraph:
                current_subgraph = None
                continue

            # 节点
            node = self.parse_node_line(line, current_subgraph)
            if node:
                nodes.append(node)
                if current_subgraph:
                    current_subgraph['nodes'].append(node['id'])
                continue

            # 边
            edge = self.parse_edge_line(line)
            if edge:
                edges.append(edge)

        return nodes, edges, subgraphs

    def parse_node_line(self, line, current_subgraph):
        patterns = [
            r'(\w+)\[([^\]]+)\]',  # A[Label]
            r'(\w+)\(([^)]+)\)',   # A(Label)
            r'(\w+)\{([^}]+)\}',   # A{Label}
            r'(\w+)',              # A
        ]
        for pattern in patterns:
            match = re.match(pattern, line)
            if match:
                node_id = match.group(1)
                label = match.group(2) if len(match.groups()) > 1 else node_id
                if '()' in pattern:
                    style = 'ellipse;whiteSpace=wrap;html=1;fillColor=#fff2cc;strokeColor=#d6b656;'
                elif '{}' in pattern:
                    style = 'rhombus;whiteSpace=wrap;html=1;fillColor=#ffe6cc;strokeColor=#d79b00;'
                else:
                    style = 'rounded=1;whiteSpace=wrap;html=1;fillColor=#dae8fc;strokeColor=#6c8ebf;'
                return {
                    'id': node_id,
                    'label': label,
                    'style': style,
                    'subgraph': current_subgraph['id'] if current_subgraph else None
                }
        return None

    def parse_edge_line(self, line):
        patterns = [
            r'(\w+)\s*-->\s*(\w+)',  # A --> B
            r'(\w+)\s*--\s*(\w+)',   # A -- B
            r'(\w+)\s*==>\s*(\w+)',  # A ==> B
            r'(\w+)\s*==\s*(\w+)',   # A == B
            r'(\w+)\s*-->\|([^|]+)\|\s*(\w+)',  # A -->|label| B
            r'(\w+)\s*--\|([^|]+)\|\s*(\w+)',   # A --|label| B
        ]
        for pattern in patterns:
            match = re.match(pattern, line)
            if match:
                groups = match.groups()
                source = groups[0]
                if len(groups) == 2:
                    target = groups[1]
                    label = ""
                else:
                    target = groups[-1]
                    label = groups[1]
                if '==' in pattern:
                    style = 'edgeStyle=orthogonalEdgeStyle;rounded=0;orthogonalLoop=1;jettySize=auto;html=1;strokeWidth=2;strokeColor=#d79b00;'
                elif '--' in pattern:
                    style = 'edgeStyle=orthogonalEdgeStyle;rounded=0;orthogonalLoop=1;jettySize=auto;html=1;strokeStyle=dashed;'
                else:
                    style = 'edgeStyle=orthogonalEdgeStyle;rounded=0;orthogonalLoop=1;jettySize=auto;html=1;'
                return {
                    'source': source,
                    'target': target,
                    'label': label,
                    'style': style
                }
        return None

    def calculate_layout(self, nodes, edges, subgraphs):
        positions = {}
        cols = 4
        for i, node in enumerate(nodes):
            row = i // cols
            col = i % cols
            x = 50 + col * 200
            y = 50 + row * 120
            positions[node['id']] = {
                'x': x,
                'y': y,
                'width': 120,
                'height': 60
            }
        return positions

    def convert_mermaid_to_drawio(self, mermaid_code, title="Diagram"):
        nodes, edges, subgraphs = self.parse_mermaid_code(mermaid_code)
        positions = self.calculate_layout(nodes, edges, subgraphs)

        # 创建 mxGraphModel（明文）
        mxgraph_model = self.create_mxgraphmodel_xml(nodes, edges, subgraphs, positions)

        # 压缩
        compressed_xml = self.compress_xml(mxgraph_model)

        # 创建最终结构
        mxfile, diagram = self.create_drawio_skeleton(title)
        diagram.text = compressed_xml  # 设置压缩内容

        return mxfile

    def save_drawio_file(self, mxfile, filename):
        rough_string = ET.tostring(mxfile, 'unicode')
        reparsed = minidom.parseString(rough_string)
        pretty_string = reparsed.toprettyxml(indent="  ")

        # 去掉 minidom 添加的多余空行
        lines = [line for line in pretty_string.splitlines() if line.strip()]
        with open(filename, 'w', encoding='utf-8') as f:
            f.write('\n'.join(lines))

        print(f"已保存: {filename}")

def extract_mermaid_from_markdown(markdown_file):
    """从Markdown文件中提取Mermaid代码块"""
    with open(markdown_file, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # 查找所有mermaid代码块
    mermaid_blocks = re.findall(r'```mermaid\s*\n(.*?)\n```', content, re.DOTALL)
    
    # 提取图表标题
    titles = re.findall(r'##\s*(图\d+[：:]\s*[^\n]+)', content)
    
    result = []
    for i, mermaid_code in enumerate(mermaid_blocks):
        title = titles[i] if i < len(titles) else f"图{i+1}"
        result.append((mermaid_code, title))
    
    return result

def main():
    """主函数"""
    converter = SimpleMermaidConverter()
    
    # 专利图文件列表
    patent_files = [
        'patents/patent_1_diagrams.md',
        'patents/patent_2_diagrams.md',
        'patents/patent_3_diagrams.md',
        'patents/patent_4_diagrams.md',
        'patents/patent_5_diagrams.md'
    ]
    
    # 创建输出目录
    output_dir = 'drawio_diagrams'
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    total_converted = 0
    
    # 处理每个专利文件
    for patent_file in patent_files:
        if not os.path.exists(patent_file):
            print(f"文件不存在: {patent_file}")
            continue
            
        print(f"处理文件: {patent_file}")
        
        # 提取Mermaid代码块
        mermaid_blocks = extract_mermaid_from_markdown(patent_file)
        
        # 获取专利名称
        patent_name = os.path.basename(patent_file).replace('_diagrams.md', '')
        
        # 转换每个图表
        for i, (mermaid_code, title) in enumerate(mermaid_blocks, 1):
            try:
                print(f"  转换图表 {i}: {title}")
                
                # 转换Mermaid到Draw.io
                mxfile = converter.convert_mermaid_to_drawio(mermaid_code, f"{patent_name}_{title}")
                
                # 保存文件
                output_filename = f"{output_dir}/{patent_name}_图{i}.drawio"
                converter.save_drawio_file(mxfile, output_filename)
                total_converted += 1
                
            except Exception as e:
                print(f"转换图表 {i} 时出错: {e}")
                import traceback
                traceback.print_exc()
    
    print(f"\n转换完成！")
    print(f"共转换了 {len(patent_files)} 个专利文件")
    print(f"成功生成 {total_converted} 个Draw.io文件")
    print(f"所有文件已保存到 {output_dir} 目录")

if __name__ == "__main__":
    main()